{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e0b2b8f-c788-4152-971d-a75c6a5f4adc",
   "metadata": {},
   "source": [
    "## 1 配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3127be26-46e8-48aa-9795-c0ee89bb5a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['http_proxy'] = 'http://127.0.0.1:1087'\n",
    "os.environ['https_proxy'] = 'http://127.0.0.1:1087'\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "\n",
    "model_name_or_path = \"/home/cc/models/asr/whisper-large-v2\"\n",
    "model_dir = \"models/whisper-large-v2-asr-int8\"\n",
    "\n",
    "language = \"Chinese (China)\"\n",
    "language_abbr = \"zh-CN\"\n",
    "language_decode = \"chinese\"\n",
    "\n",
    "task = \"transcribe\"\n",
    "dataset_name = \"mozilla-foundation/common_voice_11_0\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d569dd66-8d74-4ca7-adf3-826a16cc648d",
   "metadata": {},
   "source": [
    "## 2 导入模型和测试数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "22c7a1eb-93ab-40b9-b577-c328542418b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSpeechSeq2Seq, AutoTokenizer, AutoProcessor\n",
    "from peft import PeftConfig, PeftModel\n",
    "import torch \n",
    "\n",
    "\n",
    "peft_config = PeftConfig.from_pretrained(model_dir)\n",
    "\n",
    "base_model = AutoModelForSpeechSeq2Seq.from_pretrained(\n",
    "    peft_config.base_model_name_or_path, load_in_8bit=False, device_map=\"auto\"\n",
    ")\n",
    "\n",
    "peft_model = PeftModel.from_pretrained(base_model, model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97c0ec5d-d171-4484-bc17-eaf0dcbc6219",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# 加载前 1000 条样本\n",
    "common_voice_test = load_dataset(\n",
    "    dataset_name,\n",
    "    language_abbr,\n",
    "    split=\"test[:1000]\",  # 切片语法\n",
    "    trust_remote_code=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48a33c5d-25f5-4a06-ae45-3efd42a6ded3",
   "metadata": {},
   "outputs": [],
   "source": [
    "common_voice_test = common_voice_test.shuffle(seed=16).select(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "133944ed-8670-43c6-87da-91adf8bf8cbe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['client_id', 'path', 'audio', 'sentence', 'up_votes', 'down_votes', 'age', 'gender', 'accent', 'locale', 'segment'],\n",
       "    num_rows: 10\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "common_voice_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f438889-4033-45d4-a876-a33339c2503b",
   "metadata": {},
   "source": [
    "## 3 处理测试数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f7d2be8b-029b-4db7-b1ec-698be6fcdeb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "common_voice_test = common_voice_test.remove_columns(\n",
    "    [\"accent\", \"age\", \"client_id\", \"down_votes\", \"gender\", \"locale\", \"path\", \"segment\", \"up_votes\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "37d47025-a15b-4875-81d4-06df3e37858c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)\n",
    "processor = AutoProcessor.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)\n",
    "feature_extractor = processor.feature_extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "267db95e-a9d6-4d23-b18a-67a80769af7c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "52093bc60f6e48898e392c10c57b2dee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 降到16kHz\n",
    "common_voice_test = common_voice_test.cast_column(\"audio\", Audio(sampling_rate=16000))\n",
    "# 预处理\n",
    "def prepare_test_dataset(batch):\n",
    "    audio = batch[\"audio\"]\n",
    "    batch[\"input_features\"] = feature_extractor(audio[\"array\"], sampling_rate=audio[\"sampling_rate\"]).input_features[0]\n",
    "    return batch\n",
    "test_dataset = common_voice_test.map(prepare_test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b02dd8e0-722e-4437-8cb0-e3167c2064e7",
   "metadata": {},
   "source": [
    "## 4 评估函数\n",
    "评估指标：\n",
    " - wer: 0.5，表示词错误率（Word Error Rate）为0.5（即50%）\n",
    " - cer: 0.10256410256410256，表示字错误率（Character Error Rate）约为10.26%\n",
    " - sentence_accuracy: 0.5,表示整个句子相同的概率\n",
    " - predictions: 一个包含10个预测字符串的列表\n",
    " - references: 一个包含10个参考字符串（真实值）的列表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "93f33883-6068-404b-94e0-ab8441a5c51b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from transformers import WhisperProcessor\n",
    "import jiwer \n",
    "import torch\n",
    "import re\n",
    "\n",
    "def quick_evaluate(model, test_dataset, processor, term_list=None, batch_size=2):\n",
    "    all_predictions = []\n",
    "    all_references = []\n",
    "    term_list = term_list if term_list is not None else []  \n",
    "\n",
    "    for i in range(0, len(test_dataset), batch_size):\n",
    "        batch = test_dataset[i:i+batch_size]\n",
    "        inputs = {\"input_features\": batch[\"input_features\"]}\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            # 得到预测结果\n",
    "            generated_ids = model.generate(\n",
    "                input_features=torch.tensor(inputs[\"input_features\"]).to(model.device),\n",
    "                max_new_tokens=255\n",
    "            )\n",
    "        \n",
    "        # 解码预测结果和参考文本\n",
    "        predictions = processor.batch_decode(generated_ids, skip_special_tokens=True)\n",
    "        references = [sentence for sentence in batch[\"sentence\"]]\n",
    "        \n",
    "        all_predictions.extend(predictions)\n",
    "        all_references.extend(references)\n",
    "        print(f\"已处理 {min(i+batch_size, len(test_dataset))}/{len(test_dataset)} 条样本\")\n",
    "    \n",
    "    # 指标计算\n",
    "    wer = jiwer.wer(all_references, all_predictions)\n",
    "    cer = jiwer.cer(all_references, all_predictions)\n",
    "    correct_sentences = sum(pred == ref for pred, ref in zip(all_predictions, all_references))\n",
    "    total_sentences = len(all_references)\n",
    "    sentence_accuracy = correct_sentences / total_sentences if total_sentences > 0 else 0.0\n",
    "    \n",
    "    return {\n",
    "        \"wer\": wer,\n",
    "        \"cer\": cer,\n",
    "        \"sentence_accuracy\": sentence_accuracy,\n",
    "        \"predictions\": all_predictions,\n",
    "        \"references\": all_references\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "962918b9-34aa-463e-86d8-edcb9ef3d2e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModel(\n",
       "  (base_model): LoraModel(\n",
       "    (model): WhisperForConditionalGeneration(\n",
       "      (model): WhisperModel(\n",
       "        (encoder): WhisperEncoder(\n",
       "          (conv1): Conv1d(80, 1280, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "          (conv2): Conv1d(1280, 1280, kernel_size=(3,), stride=(2,), padding=(1,))\n",
       "          (embed_positions): Embedding(1500, 1280)\n",
       "          (layers): ModuleList(\n",
       "            (0-31): 32 x WhisperEncoderLayer(\n",
       "              (self_attn): WhisperSdpaAttention(\n",
       "                (k_proj): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (v_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=4, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=4, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (q_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=4, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=4, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "              )\n",
       "              (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (activation_fn): GELUActivation()\n",
       "              (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "              (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "              (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (decoder): WhisperDecoder(\n",
       "          (embed_tokens): Embedding(51865, 1280, padding_idx=50257)\n",
       "          (embed_positions): WhisperPositionalEmbedding(448, 1280)\n",
       "          (layers): ModuleList(\n",
       "            (0-31): 32 x WhisperDecoderLayer(\n",
       "              (self_attn): WhisperSdpaAttention(\n",
       "                (k_proj): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (v_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=4, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=4, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (q_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=4, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=4, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "              )\n",
       "              (activation_fn): GELUActivation()\n",
       "              (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (encoder_attn): WhisperSdpaAttention(\n",
       "                (k_proj): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "                (v_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=4, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=4, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (q_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=4, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=4, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (out_proj): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "              )\n",
       "              (encoder_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (fc1): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "              (fc2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "              (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "      (proj_out): Linear(in_features=1280, out_features=51865, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "peft_model = peft_model.to(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "peft_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30250402-941d-4754-91d1-e8374c4f99bd",
   "metadata": {},
   "source": [
    "## 5 评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "8c7c11de-09fb-4211-97dd-ab9b493270e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "已处理 2/10 条样本\n",
      "已处理 4/10 条样本\n",
      "已处理 6/10 条样本\n",
      "已处理 8/10 条样本\n",
      "已处理 10/10 条样本\n"
     ]
    }
   ],
   "source": [
    "evaluation_results = quick_evaluate(\n",
    "    model=peft_model,\n",
    "    test_dataset=test_dataset,\n",
    "    processor=processor\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "d9ad9178-6b5d-45ef-b8b2-bc3608cf13ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'wer': 0.5,\n",
       " 'cer': 0.10256410256410256,\n",
       " 'sentence_accuracy': 0.5,\n",
       " 'predictions': ['通过可以三人居住地到达现在居住。',\n",
       "  '曾祖父刘文。',\n",
       "  '莱斯特广场的广场。',\n",
       "  '阿布罗维亚。',\n",
       "  '齐。',\n",
       "  '正方走出家门，看见一个有刺青的流氓在家门口前在乱丢垃圾，却对他感怒不敢言。',\n",
       "  '妻子为香港行政会议成员胡红玉。',\n",
       "  '西本正的摄影工作而闻名。',\n",
       "  '它们还可以做早餐。',\n",
       "  '近几年获奖无数。'],\n",
       " 'references': ['通过科伊桑人居住地到达现在居处。',\n",
       "  '曾祖父刘文。',\n",
       "  '莱斯特广场的广场。',\n",
       "  '阿布洛维尔。',\n",
       "  '七',\n",
       "  '正方走出家门，看见一个有刺青的流氓在家门口前在乱丢垃圾，却对他敢怒不敢言。',\n",
       "  '妻子为香港行政会议成员胡红玉。',\n",
       "  '西本正的摄影工作而闻名。',\n",
       "  '他们可以做早餐',\n",
       "  '近几年获奖无数。']}"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8359974c-88e4-4921-88ec-a2b11b8e72b1",
   "metadata": {},
   "source": [
    "## 6 预测与真实结果对比"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "524b375a-2dd7-41e7-9775-12f01b998757",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "预测值：通过可以三人居住地到达现在居住。\n",
      "真实值：通过科伊桑人居住地到达现在居处。\n",
      "\n",
      "\n",
      "预测值：曾祖父刘文。\n",
      "真实值：曾祖父刘文。\n",
      "\n",
      "\n",
      "预测值：莱斯特广场的广场。\n",
      "真实值：莱斯特广场的广场。\n",
      "\n",
      "\n",
      "预测值：阿布罗维亚。\n",
      "真实值：阿布洛维尔。\n",
      "\n",
      "\n",
      "预测值：齐。\n",
      "真实值：七\n",
      "\n",
      "\n",
      "预测值：正方走出家门，看见一个有刺青的流氓在家门口前在乱丢垃圾，却对他感怒不敢言。\n",
      "真实值：正方走出家门，看见一个有刺青的流氓在家门口前在乱丢垃圾，却对他敢怒不敢言。\n",
      "\n",
      "\n",
      "预测值：妻子为香港行政会议成员胡红玉。\n",
      "真实值：妻子为香港行政会议成员胡红玉。\n",
      "\n",
      "\n",
      "预测值：西本正的摄影工作而闻名。\n",
      "真实值：西本正的摄影工作而闻名。\n",
      "\n",
      "\n",
      "预测值：它们还可以做早餐。\n",
      "真实值：他们可以做早餐\n",
      "\n",
      "\n",
      "预测值：近几年获奖无数。\n",
      "真实值：近几年获奖无数。\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for pred, ref in zip(evaluation_results['predictions'], evaluation_results['references']):\n",
    "    print('预测值：' + pred)\n",
    "    print('真实值：' + ref)\n",
    "    print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e022c8dd-8ab6-4d0a-bf49-4d20c59b3671",
   "metadata": {},
   "source": [
    "能看出有一些问题：\n",
    "专业名词识别问题（科伊桑人→可以三人）\n",
    "同音异义字处理缺陷（七→齐，敢→感）\n",
    "冗余生成倾向（添加\"还\"和句号）"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
